package com.baidubce.controller;

import java.util.Base64;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import javax.servlet.http.HttpServletRequest;

import com.baidubce.qianfan.model.image.Image2TextResponse;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestPart;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.builder.ChatBuilder;
import com.baidubce.qianfan.model.chat.ChatResponse;
import com.baidubce.qianfan.model.chat.Message;
import com.baidubce.qianfan.model.exception.ApiException;
import com.baidubce.utils.JsonUtils;
import com.baidubce.utils.SecUtils;

@RestController
public class AiChatController {
	private static final Logger logger = LoggerFactory.getLogger(AiChatController.class);
	private static final String accessKey = "qHzm2g9f2y5s7k10lXEbenZs";
	private static final String secretKey = "0NKrjr15isRenqbnHcAd0QFT3747hrS5";
	private static Qianfan qianfan = new Qianfan("OAuth", accessKey, secretKey);

	/**
	 * 发送消息接口
	 *
	 * @param params 请求参数
	 * @return 响应结果
	 */
	@PostMapping("/sendMsg")
	public String sendMsg(@RequestBody Map<String, Object> params) {
		// 验签
		if (!checkSign(params)) {
			return "签名不正确！";
		}
		String result = null;
		try {
			result = chat(String.valueOf(params.get("messages")));
		} catch (ApiException e) {
			e.printStackTrace();
			return "接口繁忙，请稍后再试！";
		} catch (Exception e) {
			e.printStackTrace();
			return "内部错误，请稍后再试！";
		}
		return result;
	}

	/**
	 * 图片上传接口
	 *
	 * @param file 图片文件
	 * @param params 请求参数
	 * @return 响应结果
	 */
	@PostMapping("/uploadImage")
	public String uploadImage(@RequestPart("image") MultipartFile file, @RequestPart("params") Map<String, Object> params) {
		// 验签
//		if (!checkSign(params)) {
//			return "签名不正确！";
//		}
		String result = null;
		try {
			String base64Image = Base64.getEncoder().encodeToString(file.getBytes());
			String prompt = String.valueOf(params.get("prompt"));
			result = recognizeImage(base64Image, prompt);
		} catch (Exception e) {
			e.printStackTrace();
			return "图片识别失败，请重试！";
		}

		return result;
	}

	/**
	 * 验签方法
	 *
	 * @param params 请求参数
	 * @return 验签结果
	 */
	private boolean checkSign(Map<String, Object> params) {
		String timestamp = String.valueOf(params.get("timestamp"));
		String messages = String.valueOf(params.get("messages"));
		String signature = String.valueOf(params.get("signature"));
		if (StringUtils.isAnyBlank(timestamp, messages, signature)) {
			return false;
		}

		// 1. 判断时间
		if ((System.currentTimeMillis() - Long.valueOf(timestamp)) > 5000) {
			// 过期
			return false;
		}
		// 2. 验签
		String p = "MYCHAT|" + timestamp + "|" + messages;
		String md5of3 = SecUtils.encoderByMd5With32Bit(SecUtils.encoderByMd5With32Bit(SecUtils.encoderByMd5With32Bit(p)));
		if (!signature.equalsIgnoreCase(md5of3)) {
			return false;
		}
		return true;
	}

	/**
	 * 调用大模型进行对话
	 *
	 * @param messages 对话记录
	 * @return 对话结果
	 * @throws Exception 如果调用过程中发生异常
	 */
	private String chat(String messages) throws Exception {
		List<Message> messageList = JsonUtils.readValues(messages, Message.class);
		messageList = messageList.stream().filter(m -> StringUtils.isNotBlank(m.getContent()))
				.collect(Collectors.toList());

		// 确保消息列表的最后一个元素是用户的消息
		if (!messageList.isEmpty() && messageList.get(messageList.size() - 1).getRole().equals("assistant")) {
			messageList.remove(messageList.size() - 1);
		}

		ChatBuilder builder = qianfan.chatCompletion()
				.model("ERNIE-Tiny-8K");

		for (Message m : messageList) {
			builder.addMessage(m);
		}
		ChatResponse response = builder.execute();
		return response.getResult();
	}

	/**
	 * 调用大模型进行图片识别
	 *
	 * @param base64Image 图片的Base64编码字符串
	 * @param prompt 提示词
	 * @return 图片识别的结果
	 * @throws Exception 如果调用过程中发生异常
	 */
	private String recognizeImage(String base64Image, String prompt) throws Exception {
		Image2TextResponse response = qianfan.image2Text()
				.model("Fuyu-8B")
				.image(base64Image)
				.prompt(prompt)
				.execute();

		return response.getResult();
	}

	public static void main(String[] args) {
		// 测试方法
//        chat("对于调休你怎么看");
		chatStream("介绍一下自己");
	}

	private static void chatStream(String message) {
		Iterator<ChatResponse> stream = qianfan.chatCompletion()
				.model("ERNIE-Tiny-8K")
				.addMessage("user", message)
				.executeStream();
		while (stream.hasNext()) {
			System.out.println(stream.next().getResult());
		}
	}
}